import os
import torch
import torchvision
from d2l import torch as d2l
import matplotlib.pyplot as plt

n = 5
voc_dir = d2l.download_extract('voc2012', os.path.join(
        'VOCdevkit', 'VOC2012'))
train_features,train_labels = d2l.read_voc_images(voc_dir,True)
imgs = train_features[0:n] + train_labels[0:n]
imgs = [img.permute(1,2,0) for img in imgs]
d2l.show_images(imgs,2,n)
plt.show()

batch_size = 256
crop_size = (320,480)
train_iter,test_iter = d2l.load_data_voc(batch_size,crop_size)
print(train_iter[0].shape)
print(test_iter[0].shape)
